package com.sunzm.hbase;

import com.alibaba.fastjson.JSONObject;
import com.sunzm.common.utils.ParameterTool;
import org.apache.commons.io.FileUtils;
import org.apache.commons.io.LineIterator;
import org.apache.commons.lang3.ArrayUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.*;
import org.apache.hadoop.hbase.client.*;
import org.apache.hadoop.hbase.filter.ColumnValueFilter;
import org.apache.hadoop.hbase.util.Bytes;

import java.io.File;
import java.io.IOException;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

/**
 * 从HBase中查询数据，生成文件
 *
 * @author Administrator
 * @version 1.0
 * @date 2021-06-10 16:17
 */
public class HBaseDataUtils {
    private Configuration conf = null;
    private String hbaseZookeeperQuorum = null;
    private Connection connection = null;

    public void init() {

        System.out.println(HConstants.ZOOKEEPER_QUORUM + ":" + hbaseZookeeperQuorum);

        conf = HBaseConfiguration.create();

        System.out.println("初始化HBase连接.....");
        conf.set("hbase.zookeeper.quorum", hbaseZookeeperQuorum);

        try {
            connection = ConnectionFactory.createConnection(conf);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public static void main(String[] args) {

        ParameterTool params = ParameterTool.fromArgs(args);

        String hbaseZookeeperQuorumStr = params.getRequired(HConstants.ZOOKEEPER_QUORUM);

        String tableName = params.getRequired("tableName");

        /**
         * get 需要rowkey
         * scan 需要startKey和stopkey
         * toFile 需要readFileName, writeFileName，每一行必须是一个rowkey
         */
        String queryType = params.getRequired("queryType");
        //需要查询的列，默认情况下，查询所有列
        //cid, companyId
        String queryColumn = params.get("queryColumn", "ALL");
        //需要过滤的列名：只支持根据一个列过滤,列名必须带上列簇名，例如: info:senderType
        String filterColumnName = params.get("filterColumnName", "").trim();
        //需要过滤的列值：只支持根据一个列过滤
        String filterColumnValue = params.get("filterColumnValue", "").trim();

        String readFileName = null;
        String writeFileName = null;
        String rowkey = null;
        String startKey = null;
        String stopkey = null;

        HBaseDataUtils hBaseDataUtils = new HBaseDataUtils();
        hBaseDataUtils.hbaseZookeeperQuorum = hbaseZookeeperQuorumStr;

        switch (queryType) {
            case "toFile": {

                readFileName = params.get("readFileName", "rowkey.dat");
                writeFileName = params.get("writeFileName", "queryResult.dat");

                try {
                    hBaseDataUtils.hbaseData2File(tableName, readFileName, writeFileName, queryColumn,
                            filterColumnName, filterColumnValue);
                } catch (IOException e) {
                    e.printStackTrace();
                }

                break;
            }
            default: {
                System.out.println("暂时不支持的 queryType: " + queryType);
            }
        }

        System.out.println("任务执行完成!");

        System.exit(0);
    }

    private void hbaseData2File(String tableNameStr, String readFileName, String writeFileName,
                                String queryColumn, String filterColumnName, String filterColumnValue) throws IOException {
        init();

        TableName tableName = TableName.valueOf(tableNameStr);
        Table table = connection.getTable(tableName);

        Scan scan = new Scan();

        scan.setAttribute(Scan.SCAN_ATTRIBUTES_TABLE_NAME, tableNameStr.getBytes());

        scan.setCaching(2000);

        scan.setCacheBlocks(false);

        if(StringUtils.isNotBlank(filterColumnName) && StringUtils.isNotBlank(filterColumnValue)){
            //设置过滤器 info:senderType
            String realFamilyName = "info";
            String realFilterColumnName = filterColumnName;
            if(StringUtils.contains(filterColumnName, ":")){
                String[] filterColumnNameArray = StringUtils.splitByWholeSeparatorPreserveAllTokens(filterColumnName, ":");

                realFamilyName = filterColumnNameArray[0];
                realFilterColumnName = filterColumnNameArray[1];
            }
            ColumnValueFilter filter = new ColumnValueFilter(Bytes.toBytes(realFamilyName), Bytes.toBytes(realFilterColumnName),
                    CompareOperator.EQUAL, Bytes.toBytes(filterColumnValue));
            scan.setFilter(filter);
        }

        //在当前目录生成文件
        File readFile = new File(readFileName);
        File writeFile = new File(writeFileName);
        //File zipedFile = new File(writeFileName + ".zip");

        //删除文件
        FileUtils.deleteQuietly(writeFile);
        //FileUtils.deleteQuietly(zipedFile);

        LineIterator lineIterator = FileUtils.lineIterator(readFile, "UTF-8");


        JSONObject jSONObject = new JSONObject();

        List<String> lines = new ArrayList<>();

        while (lineIterator.hasNext()) {
            String line = lineIterator.nextLine();
            String rowKey = StringUtils.trim(line);
            String startRowKey = rowKey;
            String stopRowKey = rowKey + "~";

            scan.withStartRow(startRowKey.getBytes("UTF-8"));
            scan.withStopRow(stopRowKey.getBytes("UTF-8"));

            ResultScanner scanner = table.getScanner(scan);

            Iterator<Result> iterator = scanner.iterator();

            while (iterator.hasNext()) {
                Result result = iterator.next();

                CellScanner cellScanner = result.cellScanner();

                jSONObject.clear();
                while (cellScanner.advance()) {
                    Cell cell = cellScanner.current();

                    byte[] qualifierBytes = cell.getQualifierArray();
                    byte[] valueBytes = cell.getValueArray();

                    String qualifierStr = new String(qualifierBytes, cell.getQualifierOffset(), cell.getQualifierLength());
                    String valueStr = new String(valueBytes, cell.getValueOffset(), cell.getValueLength());

                    if (StringUtils.equals("ALL", queryColumn)) {
                        //如果查询所有列
                        jSONObject.put(qualifierStr, valueStr);
                    } else {
                        String[] queryColumnArray = StringUtils.splitByWholeSeparatorPreserveAllTokens(queryColumn, ",");

                        if (ArrayUtils.contains(queryColumnArray, qualifierStr)) {
                            //只保存需要的列
                            jSONObject.put(qualifierStr, valueStr);
                        }
                    }
                }

                lines.add(jSONObject.toJSONString());

                if (lines.size() >= 10000) {
                    System.out.println("达到10000条,开始写文件...");
                    FileUtils.writeLines(writeFile, "UTF-8", lines, System.lineSeparator(), true);
                    //写文件完成后，清空ArrayList
                    lines.clear();
                }
            }


        }


        if (lines.size() > 0) {
            System.out.println("最好一次写文件, 数据条数: " + lines.size());
            FileUtils.writeLines(writeFile, "UTF-8", lines, System.lineSeparator(), true);
            lines.clear();
        }

        close(connection);

        //压缩文件
       /* try {
            //FileOpeUtils.zipFile(writeFile, zipedFile);
        } catch (Exception e) {
            System.err.println("文件压缩异常...." + e.getMessage());
        }*/
    }

    private String getByRowkey(String tableNameStr, String rowkey) throws IOException {
        init();

        TableName tableName = TableName.valueOf(tableNameStr);
        Table table = connection.getTable(tableName);

        Get get = new Get(rowkey.getBytes(Charset.forName("UTF-8")));

        Result result = table.get(get);

        CellScanner cellScanner = result.cellScanner();

        JSONObject jSONObject = new JSONObject();
        while (cellScanner.advance()) {
            Cell cell = cellScanner.current();

            byte[] qualifierBytes = cell.getQualifierArray();
            byte[] valueBytes = cell.getValueArray();

            String qualifierStr = new String(qualifierBytes, cell.getQualifierOffset(), cell.getQualifierLength());
            String valueStr = new String(valueBytes, cell.getValueOffset(), cell.getValueLength());

            jSONObject.put(qualifierStr, valueStr);
        }
        String res = "";
        if (!jSONObject.isEmpty()) {
            res = jSONObject.toJSONString();
        }

        close(connection);

        return res;
    }

    public void close(Connection connection) {
        System.out.println("关闭HBase连接.....");
        if (connection != null && !connection.isClosed()) {
            try {
                connection.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }
}
